from datetime import datetime

import matplotlib.pyplot as plt
import networkx as nx
import numpy as np


def saveGraph(bom_structure):
    # 初始化有向图
    G = nx.DiGraph()

    # tuple/str
    def addNode(data):
        if isinstance(data, tuple):
            part_name, part_weight = data
            G.add_node(part_name, weight=part_weight)
        else:
            G.add_node(data)

    def addEdge(source, target):
        if isinstance(target, str):
            G.add_edge(source, target)
        elif isinstance(target, tuple):  # 带权重的边
            G.add_edge(source, target[0], weight=target[1])

    # 规范化数据并构建有向图
    for key, value in bom_structure.items():
        # 添加节点
        addNode(key)

        # 检查值是否是字典
        if isinstance(value, dict):
            for child, connections in value.items():
                # 添加节点
                addNode(child)

                # 检查连接是否是集合
                if isinstance(connections, set):
                    for target in connections:
                        addEdge(child, target)
                # 检查连接是否是字典
                elif isinstance(connections, dict):
                    for target, weight in connections.items():
                        addEdge(child, target)

        # 添加父节点到子节点的边
        for child in value:
            if isinstance(value[child], set) or isinstance(value[child], dict):
                addEdge(key, child)

    # 完成图模型的构建
    print(G)

    # 动态调整显示大小
    node_count = len(G.nodes())
    node_size = 500 * (1 / node_count ** 0.5)  # 根据节点数量调整节点大小

    # 使用kamada_kawai_layout布局算法
    pos = nx.kamada_kawai_layout(G)

    # 绘制有向图
    plt.figure(figsize=(12, 12))
    node_colors = ['skyblue' for _ in range(len(G.nodes))]  # 可以基于权重动态分配颜色
    edge_colors = ['gray' for _ in range(len(G.edges))]  # 边的颜色
    nx.draw(G, pos, with_labels=True, node_size=node_size, font_size=7, arrows=True, node_color=node_colors,
            edge_color=edge_colors, connectionstyle="arc3")

    # 绘制边的权重，避免标签重叠
    edge_labels = nx.get_edge_attributes(G, 'weight')
    for edge, weight in edge_labels.items():
        if weight is not None:
            # 计算边的中点坐标
            x0, y0 = pos[edge[0]]
            x1, y1 = pos[edge[1]]
            midpoint = ((x0 + x1) / 2, (y0 + y1) / 2)

            # 为标签添加一些随机偏移量以避免重叠
            offset = 0.01
            xytext = (midpoint[0] + (x1 - x0) * offset, midpoint[1] + (y1 - y0) * offset)

            # 在边的中间位置绘制权重标签，并使用红色字体
            plt.annotate(str(weight), xy=(x0, y0), xytext=xytext, textcoords='data',
                         arrowprops=dict(arrowstyle="->", connectionstyle="arc3", color='red'),
                         color='red', fontsize=8, ha='center', va='center')

    plt.title('BOM Structure Graph')
    plt.axis('off')
    plt.savefig(f'BOM_Graph_{datetime.now().strftime("%Y%m%d_%H%M%S")}.png', dpi=300, bbox_inches='tight')
    # plt.show()


def saveAttrFeature(bom_attr):
    # 1. 遍历数据，提取唯一的属性
    unique_attributes = set()
    for attr_info in bom_attr.values():
        unique_attributes.update(attr_info.keys())

    # 2. 构建属性值到索引的映射字典
    attribute_to_index = {attr: i for i, attr in enumerate(unique_attributes)}

    # 3. 对属性信息进行独热编码
    feature_matrix = np.zeros((len(bom_attr), len(unique_attributes)), dtype=int)
    for i, attr_info in enumerate(bom_attr.values()):
        for attr, value in attr_info.items():
            j = attribute_to_index[attr]
            feature_matrix[i, j] = 1

    # 设置中文字体
    plt.rcParams['font.sans-serif'] = ['SimHei']
    plt.rcParams['axes.unicode_minus'] = False

    # 绘制特征矩阵图
    plt.figure()
    plt.imshow(feature_matrix, cmap='binary', interpolation='nearest')
    plt.xticks(np.arange(len(unique_attributes)), unique_attributes, rotation=45)
    plt.yticks(np.arange(len(bom_attr)), bom_attr.keys())
    plt.xlabel('Attr')
    plt.ylabel('SBB')
    plt.title('BOM Attr Feature Matrix')
    plt.savefig(f'BOM_AttrFeature_{datetime.now().strftime("%Y%m%d_%H%M%S")}.png', dpi=300, bbox_inches='tight')
    # plt.show()


if __name__ == '__main__':
    bom_structure = {
        'WGPart1': {
            'SBB3': {('SBB8', 2), 'SBB4', 'SBB5'},
            'SBB4': {'SBB3': {('SBB8', 2)}, ('SBB9', 2): {}},
            'SBB5': {'SBB1', 'SBB6'},
            'SBB6': {'SBB5', ('SBB3', 10)},
            ('SBB8', 3): {('SBB1', 10)}
        },
        ('WGPart2', 2): {},
        'WGPart3': {'SBB1': {'SBB2'}},
        'WGPart4': {'SBB7': {'SBB8', ('SBB9', 2)}}
    }
    bom_attr = {
        "SBB1": {"大小": "10寸", "材质": "玻璃", "重量": "200g"},
        "SBB2": {"型号": "ABC123", "功耗": "50W"},
        "SBB3": {"品牌": "XYZ", "尺寸": "5mm"},
        "SBB4": {"类型": "机械键盘", "颜色": "黑色", "重量": "300g"},
        "SBB5": {"型号": "Intel Core i7", "功耗": "65W"}
    }

    saveGraph(bom_structure)
    saveAttrFeature(bom_attr)
